#!/usr/bin/env bash

set -e
set -o pipefail

log() {
    local fname=${BASH_SOURCE[1]##*/}
    echo -e "$(date '+%Y-%m-%dT%H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
}
SECONDS=0

# general configuration
stage=1
stop_stage=100
nj=32
set_name=train
set_type=far
rttm_file=
rttm_postfix=

log "$0 $*"
. utils/parse_options.sh

. ./cmd.sh || exit 1
. ./path.sh || exit 1

pip show -f nara_wpe >/dev/null || pip install nara_wpe

if [ -z $BEAMFORMIT ]; then
    export BEAMFORMIT=$KALDI_ROOT/tools/BeamformIt
fi
export PATH=${PATH}:$BEAMFORMIT
! hash BeamformIt && echo "Missing BeamformIt, run 'cd ../../../tools/kaldi/tools; extras/install_beamformit.sh; cd -;'" && exit 1

if [ $# != 2 ]; then
    log "Usage: $0 <corpus-data-dir> <enhancement-dir>"
    log " $0 /path/misp2021 /path/wpe_output"
    exit 1
fi

audio_dir=$1
output_dir=$2

log "Start Speech Enhancement."

# WPE
if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
    log "stage 1: Start WPE for ${set_name}_${set_type}."
    python local/find_wav.py \
        --audio-dir ${audio_dir} \
        --scp-dir ${output_dir}/log \
        --scp-name wpe_${set_name}_${set_type} \
        --set-type ${set_type} \
        --nj ${nj}
    for n in `seq ${nj}`; do
        cat <<-EOF >${output_dir}/log/wpe_${set_name}_${set_type}.${n}.sh
python local/run_wpe.py \
    --wav-scp ${output_dir}/log/wpe_${set_name}_${set_type}.${n}.scp \
    --audio-dir ${audio_dir} \
    --output-dir ${output_dir}
EOF
    done
    chmod a+x ${output_dir}/log/wpe_${set_name}_${set_type}.*.sh
    ${train_cmd} JOB=1:${nj} ${output_dir}/log/wpe_${set_name}_${set_type}.JOB.log \
        ${output_dir}/log/wpe_${set_name}_${set_type}.JOB.sh
fi

# BeamformIt
if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
    log "stage 2: Start BeamformIt for ${set_name}_${set_type}"
    python local/find_wav.py \
        --audio-dir ${audio_dir} \
        --scp-dir ${output_dir}/log \
        --scp-name beamformit_${set_name}_${set_type} \
        --set-type ${set_type}
    python local/run_beamformit.py \
        --beamformit-tool ${BEAMFORMIT}/BeamformIt \
        --config-file conf/beamformit.cfg \
        --source-dir / \
        --channel-scp ${output_dir}/log/beamformit_${set_name}_${set_type}.scp
fi

# GSS
if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
    log "stage 3: Start GSS for ${set_name}_${set_type}"
    if [ ! -d pb_chime5/ ]; then
        log "Please install pb_chime5 by local/install_pb_chime5.sh"
        exit 1
    fi

    python3 local/find_wav.py \
        --audio-dir ${audio_dir} \
        --scp-dir ${output_dir}/log \
        --scp-name gss_${set_name}_${set_type} \
        --set-type ${set_type} \
        --nj ${nj}

    for n in `seq ${nj}`; do
        cat <<-EOF >${output_dir}/log/gss_${set_name}_${set_type}.${n}.sh
python3 local/run_gss.py \
    --wav-scp ${output_dir}/log/gss_${set_name}_${set_type}.${n}.scp \
    --audio-dir ${audio_dir} \
    --output-dir ${output_dir} \
    --set-type ${set_type}
EOF
    done
    chmod a+x ${output_dir}/log/gss_${set_name}_${set_type}.*.sh
    ${train_cmd} JOB=1:${nj} ${output_dir}/log/gss_${set_name}_${set_type}.JOB.log \
        ${output_dir}/log/gss_${set_name}_${set_type}.JOB.sh
fi

if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
    log "stage 4: Start GSS suing RTTM file for ${set_name}_${set_type}"

    python local/find_wav.py \
        --audio-dir ${output_dir} \
        --scp-dir ${output_dir}/log \
        --scp-name gss_rttm_${set_name}_${set_type} \
        --set-type ${set_type} \
        --nj ${nj}

    for n in `seq ${nj}`; do
        cat <<-EOF >${output_dir}/log/gss_${set_name}_${set_type}_rttm_${rttm_postfix}.${n}.sh
python local/run_gss_rttm.py \
    --wav-scp ${output_dir}/log/gss_rttm_${set_name}_${set_type}.${n}.scp \
    --audio-dir ${output_dir} \
    --rttm-file ${rttm_file} \
    --output-dir ${output_dir} \
    --set-type ${set_type} \
    --rttm-postfix ${rttm_postfix}
EOF
    done
    chmod a+x ${output_dir}/log/gss_${set_name}_${set_type}_rttm_${rttm_postfix}.*.sh
    ${train_cmd} JOB=1:${nj} ${output_dir}/log/gss_${set_name}_${set_type}_rttm_${rttm_postfix}.JOB.log \
        ${output_dir}/log/gss_${set_name}_${set_type}_rttm_${rttm_postfix}.JOB.sh
fi

log "End speech enhancement"
